import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os

def extract_images_from_bag(bag_file, output_folder, topic_name):
    bridge = CvBridge()
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    print(f"正在打开 bag 文件: {bag_file}")
    with rosbag.Bag(bag_file, 'r') as bag:
        count = 0
        for topic, msg, t in bag.read_messages(topics=[topic_name]):
            if topic == topic_name:
                try:
                    # 将 ROS Image 消息转换为 OpenCV 图像
                    cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
                    # 保存图像文件
                    image_filename = os.path.join(output_folder, f"frame_{count:05d}.png")
                    cv2.imwrite(image_filename, cv_image)
                    print(f"保存图像: {image_filename}")
                    count += 1
                except Exception as e:
                    print(f"处理消息时发生错误: {e}")
    print(f"图像提取完成。共提取了 {count} 张图像到 {output_folder}")

if __name__ == '__main__':
    # 替换为你的 bag 文件路径和输出文件夹
    bag_file_path = os.path.expanduser("/home/youngqi/Desktop/2025-06-24-21-53-23.bag") # 假设文件在 home 目录下
    output_images_folder = "extracted_frames"
    image_topic = "image_view/image_view"

    extract_images_from_bag(bag_file_path, output_images_folder, image_topic)